import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                           QHBoxLayout, QLabel, QPushButton, QGroupBox, 
                           QSpinBox, QDoubleSpinBox, QComboBox, QCheckBox)
from PyQt5.QtCore import Qt, QTimer
from PyQt5.QtGui import QImage, QPixmap
import pyrealsense2 as rs
import open3d as o3d

class RealSenseViewer(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("RealSense Viewer")
        self.setGeometry(100, 100, 1280, 720)
        
        # 初始化RealSense
        self.pipeline = rs.pipeline()
        self.config = rs.config()
        
        # 设置流配置
        self.config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
        self.config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
        self.config.enable_stream(rs.stream.infrared, 1, 640, 480, rs.format.y8, 30)
        
        # 启动流
        self.pipeline.start(self.config)
        
        # 创建主窗口布局
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        self.main_layout = QHBoxLayout(self.central_widget)
        
        # 创建左侧图像显示区域
        self.setup_image_display()
        
        # 创建右侧控制面板
        self.setup_control_panel()
        
        # 设置定时器更新图像
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frames)
        self.timer.start(33)  # 约30fps
        
    def setup_image_display(self):
        """设置图像显示区域"""
        image_layout = QVBoxLayout()
        
        # 创建三个图像显示区
        self.color_label = QLabel()
        self.depth_label = QLabel()
        self.ir_label = QLabel()
        
        # 设置固定大小
        for label in [self.color_label, self.depth_label, self.ir_label]:
            label.setFixedSize(640, 360)
            label.setStyleSheet("border: 1px solid gray")
        
        image_layout.addWidget(self.color_label)
        image_layout.addWidget(self.depth_label)
        image_layout.addWidget(self.ir_label)
        
        self.main_layout.addLayout(image_layout)
        
    def setup_control_panel(self):
        """设置控制面板"""
        control_layout = QVBoxLayout()
        
        # 点云控制组
        pointcloud_group = QGroupBox("点云控制")
        pointcloud_layout = QVBoxLayout()
        
        # 显示点云按钮
        self.show_pointcloud_button = QPushButton("显示点云")
        self.show_pointcloud_button.clicked.connect(self.show_pointcloud)
        pointcloud_layout.addWidget(self.show_pointcloud_button)
        
        # 保存点云按钮
        save_pointcloud_button = QPushButton("保存点云")
        save_pointcloud_button.clicked.connect(self.save_pointcloud)
        pointcloud_layout.addWidget(save_pointcloud_button)
        
        pointcloud_group.setLayout(pointcloud_layout)
        control_layout.addWidget(pointcloud_group)
        
        # 相机控制组
        camera_group = QGroupBox("相机控制")
        camera_layout = QVBoxLayout()
        
        # 分辨率选择
        resolution_layout = QHBoxLayout()
        resolution_layout.addWidget(QLabel("分辨率:"))
        resolution_combo = QComboBox()
        resolution_combo.addItems(["640x480", "1280x720", "1920x1080"])
        resolution_layout.addWidget(resolution_combo)
        camera_layout.addLayout(resolution_layout)
        
        # 帧率控制
        fps_layout = QHBoxLayout()
        fps_layout.addWidget(QLabel("帧率:"))
        fps_spin = QSpinBox()
        fps_spin.setRange(6, 90)
        fps_spin.setValue(30)
        fps_layout.addWidget(fps_spin)
        camera_layout.addLayout(fps_layout)
        
        camera_group.setLayout(camera_layout)
        control_layout.addWidget(camera_group)
        
        # 深度控制组
        depth_group = QGroupBox("深度控制")
        depth_layout = QVBoxLayout()
        
        # 深度范围控制
        depth_range_layout = QHBoxLayout()
        depth_range_layout.addWidget(QLabel("深度范围:"))
        min_depth = QDoubleSpinBox()
        max_depth = QDoubleSpinBox()
        min_depth.setRange(0, 10)
        max_depth.setRange(0, 10)
        min_depth.setValue(0.1)
        max_depth.setValue(10.0)
        depth_range_layout.addWidget(min_depth)
        depth_range_layout.addWidget(max_depth)
        depth_layout.addLayout(depth_range_layout)
        
        depth_group.setLayout(depth_layout)
        control_layout.addWidget(depth_group)
        
        # 滤波器控制组
        filter_group = QGroupBox("滤波器")
        filter_layout = QVBoxLayout()
        
        # 空间滤波器
        spatial_check = QCheckBox("空间滤波器")
        filter_layout.addWidget(spatial_check)
        
        # 时间滤波器
        temporal_check = QCheckBox("时间滤波器")
        filter_layout.addWidget(temporal_check)
        
        filter_group.setLayout(filter_layout)
        control_layout.addWidget(filter_group)
        
        # 录制控制组
        record_group = QGroupBox("录制控制")
        record_layout = QVBoxLayout()
        
        # 录制按钮
        record_button = QPushButton("开始录制")
        record_layout.addWidget(record_button)
        
        # 保存按钮
        save_button = QPushButton("保存图像")
        record_layout.addWidget(save_button)
        
        record_group.setLayout(record_layout)
        control_layout.addWidget(record_group)
        
        # 添加弹性空间
        control_layout.addStretch()
        
        self.main_layout.addLayout(control_layout)
        
    def update_frames(self):
        """更新图像帧"""
        try:
            # 等待新的帧
            frames = self.pipeline.wait_for_frames()
            
            # 获取各个流的帧
            color_frame = frames.get_color_frame()
            depth_frame = frames.get_depth_frame()
            ir_frame = frames.get_infrared_frame(1)
            
            if not color_frame or not depth_frame or not ir_frame:
                return
            
            # 转换为numpy数组
            color_image = np.asanyarray(color_frame.get_data())
            depth_image = np.asanyarray(depth_frame.get_data())
            ir_image = np.asanyarray(ir_frame.get_data())
            
            # 将深度图转换为伪彩色图像
            depth_colormap = cv2.applyColorMap(
                cv2.convertScaleAbs(depth_image, alpha=0.03), 
                cv2.COLORMAP_JET
            )
            
            # 转换为QImage并显示
            self.display_image(color_image, self.color_label)
            self.display_image(depth_colormap, self.depth_label)
            self.display_image(ir_image, self.ir_label, True)
            
        except Exception as e:
            print(f"更新帧时出错: {str(e)}")
            
    def display_image(self, image, label, is_grayscale=False):
        """将numpy数组转换为QImage并显示在QLabel上"""
        if is_grayscale:
            h, w = image.shape
            bytes_per_line = w
            q_image = QImage(image.data, w, h, bytes_per_line, QImage.Format_Grayscale8)
        else:
            h, w, ch = image.shape
            bytes_per_line = ch * w
            q_image = QImage(image.data, w, h, bytes_per_line, QImage.Format_BGR888)
            
        pixmap = QPixmap.fromImage(q_image)
        label.setPixmap(pixmap.scaled(label.size(), Qt.KeepAspectRatio))
        
    def show_pointcloud(self):
        """显示点云"""
        try:
            # 等待新的帧
            frames = self.pipeline.wait_for_frames()
            
            # 获取深度帧和彩色帧
            depth_frame = frames.get_depth_frame()
            color_frame = frames.get_color_frame()
            
            if not depth_frame or not color_frame:
                return
                
            # 创建点云
            pc = rs.pointcloud()
            pc.map_to(color_frame)
            points = pc.calculate(depth_frame)
            
            # 获取顶点数据
            vertices = np.asanyarray(points.get_vertices())
            texcoords = np.asanyarray(points.get_texture_coordinates())
            
            # 获取颜色数据
            color_image = np.asanyarray(color_frame.get_data())
            
            # 创建Open3D点云对象
            pcd = o3d.geometry.PointCloud()
            
            # 设置点云坐标
            pcd.points = o3d.utility.Vector3dVector(vertices.reshape(-1, 3))
            
            # 设置点云颜色
            colors = []
            for texcoord in texcoords:
                u = int(texcoord[0] * color_image.shape[1])
                v = int(texcoord[1] * color_image.shape[0])
                if 0 <= u < color_image.shape[1] and 0 <= v < color_image.shape[0]:
                    color = color_image[v, u] / 255.0
                    colors.append(color[::-1])  # BGR to RGB
                else:
                    colors.append([0, 0, 0])
                    
            pcd.colors = o3d.utility.Vector3dVector(np.asarray(colors))
            
            # 显示点云
            o3d.visualization.draw_geometries([pcd])
            
        except Exception as e:
            print(f"显示点云时出错: {str(e)}")
            
    def save_pointcloud(self):
        """保存点云"""
        try:
            # 等待新的帧
            frames = self.pipeline.wait_for_frames()
            
            # 获取深度帧和彩色帧
            depth_frame = frames.get_depth_frame()
            color_frame = frames.get_color_frame()
            
            if not depth_frame or not color_frame:
                return
                
            # 创建点云
            pc = rs.pointcloud()
            pc.map_to(color_frame)
            points = pc.calculate(depth_frame)
            
            # 保存为PLY文件
            points.export_to_ply("pointcloud.ply", color_frame)
            print("点云已保存为 pointcloud.ply")
            
        except Exception as e:
            print(f"保存点云时出错: {str(e)}")
            
    def closeEvent(self, event):
        """关闭事件处理"""
        self.pipeline.stop()
        event.accept()

if __name__ == '__main__':
    app = QApplication(sys.argv)
    viewer = RealSenseViewer()
    viewer.show()
    sys.exit(app.exec_()) 